{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import lasagne\n",
    "import time\n",
    "from nbfinder import NotebookFinder\n",
    "import sys\n",
    "sys.meta_path.append(NotebookFinder())\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "from matplotlib import patches\n",
    "\n",
    "from helper_fxns import *\n",
    "\n",
    "#import data_loader\n",
    "#from data_loader import load_classification_dataset, load_detection_dataset\n",
    "#from build_hur_detection_network import build_det_network\n",
    "#from build_hur_classif_network import build_classif_network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Plotter(object):\n",
    "    def __init__(self, kwargs, fns, max_ims, iterator):\n",
    "        self.kwargs = kwargs\n",
    "        self.fns = fns\n",
    "        self.max_ims = max_ims\n",
    "        self.iterator = iterator\n",
    "        self.epoch = 0\n",
    "    def get_encoder_fmaps(self, test_kwargs, tr_kwargs):\n",
    "        fmaps_tr_dir = join(self.kwargs[\"save_path\"], \"tr_fmaps\")\n",
    "        fmaps_te_dir = join(self.kwargs[\"save_path\"], \"te_fmaps\")\n",
    "        makedir_if_not_there(fmaps_tr_dir)\n",
    "        makedir_if_not_there(fmaps_te_dir)\n",
    "        i =0\n",
    "        for i, (x,y) in enumerate(self.iterator(**test_kwargs).iterate()):\n",
    "            fmap = self.fns[\"hid\"](x)\n",
    "            pickle.dump([x,y,fmap],open(join(fmaps_te_dir, \"te_fmaps\" + str(i) + \".pkl\"), \"w\"))\n",
    "            if i + 1 > self.max_ims:\n",
    "                break\n",
    "\n",
    "\n",
    "        for i, (x,y) in enumerate(self.iterator(**tr_kwargs).iterate()):\n",
    "            fmap = self.fns[\"hid\"](x)\n",
    "            pickle.dump([x,y,fmap],open(join(fmaps_tr_dir, \"tr_fmaps\" + str(i) + \".pkl\"), \"w\"))\n",
    "            if i + 1 > self.max_ims:\n",
    "                break\n",
    "        \n",
    "            \n",
    "        \n",
    "    def postproc_ims(self, test_kwargs):\n",
    "        j =0 \n",
    "        for i, (x,y) in enumerate(self.iterator(**test_kwargs).iterate()):\n",
    "            if i%4==0:\n",
    "                j+=1\n",
    "                self.plot_ims(x,y,\"test\", j)\n",
    "        \n",
    "\n",
    "        \n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "    def plot_learn_curve(self, metrics):\n",
    "        for k in metrics.keys():\n",
    "            if \"time\" not in k:\n",
    "                k = k.replace(\"tr_\", \"\")\n",
    "                k = k.replace(\"val_\", \"\")\n",
    "                self._plot_learn_curve(k, metrics)\n",
    "        \n",
    "    def _plot_learn_curve(self,metric, metrics):\n",
    "        plt.clf()\n",
    "        plt.figure(1)\n",
    "        plt.clf()\n",
    "        ax = plt.subplot(111)\n",
    "        ax.set_title('Train/Val %s' %(metric))\n",
    "        ax.plot(metrics['tr_' + metric], label='train ' + metric)\n",
    "        ax.plot(metrics['val_' + metric], label='val ' + metric)\n",
    "        # Shrink current axis's height by 10% on the bottom\n",
    "        box = ax.get_position()\n",
    "        ax.set_position([box.x0, box.y0 + box.height * 0.1,\n",
    "                         box.width, box.height * 0.9])\n",
    "\n",
    "        # Put a legend below current axis\n",
    "        ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.05),\n",
    "                  fancybox=True, shadow=True, ncol=5)\n",
    "\n",
    "\n",
    "        curves_path = join(self.kwargs['save_path'], \"learn_curves\")\n",
    "        makedir_if_not_there(curves_path)\n",
    "        plt.savefig(\"%s/%s_learning_curve.png\"%(curves_path,metric))\n",
    "        #plt.show()\n",
    "        plt.clf()\n",
    "        \n",
    "    \n",
    "    def do_plots(self,x,y,type_, num, epoch):\n",
    "        self.epoch = epoch\n",
    "        if num >= self.max_ims:\n",
    "            return\n",
    "        if self.epoch % 2 == 0:\n",
    "            if self.kwargs[\"ignore_plot_fails\"]:\n",
    "                try:\n",
    "                    self.plot_ims(x, y, type_, num)\n",
    "                except:\n",
    "                    pass\n",
    "            else:\n",
    "                self.plot_ims(x, y, type_, num)\n",
    "    \n",
    "    def plot_ims(self, x, y, type_, num):\n",
    "\n",
    "\n",
    "\n",
    "        pred_boxes, gt_boxes = self.fns['box'](x,y)\n",
    "        self._plot_im_with_boxes(x,pred_boxes, gt_boxes, num=num, name=type_)\n",
    "        if self.kwargs['lambda_ae'] != 0:\n",
    "            xrec = self.fns['rec'](x)\n",
    "            self._plot_reconstructed_ims(x,xrec,num=num, name=type_)\n",
    "\n",
    "\n",
    "        \n",
    "       \n",
    "    def _plot_im_with_boxes(self,ims,pred_boxes, gt_boxes, num=1, name=\"tr\"):\n",
    "        \"\"\"ims is a N,vars,X,y tensor\n",
    "            pred_boxes is a list of dicts where each value of dict is a list of bbox tuples\"\"\"\n",
    "        if self.kwargs[\"3D\"]:\n",
    "            n_ims = ims.shape[2]\n",
    "        else:\n",
    "            n_ims = ims.shape[0]\n",
    "\n",
    "#         print len(pred_boxes)\n",
    "#         print len(gt_boxes)\n",
    "#         print len(pred_boxes[0])\n",
    "#         print len(gt_boxes[0])\n",
    "        channels = 1 #ims.shape[1]\n",
    "        tmq_ind = 6\n",
    "        if self.kwargs[\"3D\"]:\n",
    "            plt.figure(3, figsize=(50,50))\n",
    "        else:\n",
    "            plt.figure(3, figsize=(7,7))\n",
    "        plt.clf()\n",
    "        \n",
    "\n",
    "        count=0\n",
    "        for i in range(n_ims):\n",
    "            if self.kwargs[\"3D\"]:\n",
    "                if i %2 != 0:\n",
    "                    continue\n",
    "            for j in range(channels):\n",
    "                count+= 1\n",
    "                sp = plt.subplot(n_ims,channels, count)\n",
    "                if self.kwargs[\"3D\"]:\n",
    "                    sp.imshow(ims[0,tmq_ind,i][::-1])\n",
    "                else:\n",
    "                    sp.imshow(ims[i,tmq_ind][::-1])\n",
    "\n",
    "                if self.kwargs[\"3D\"]:\n",
    "                    i = i / 2\n",
    "                    \n",
    "                b = 1\n",
    "                for b,box in enumerate(gt_boxes[i]):\n",
    "                    self.add_bbox(sp, box, color='g')\n",
    "                    \n",
    "                  \n",
    "                # top two most confident\n",
    "                ind = 3*len(gt_boxes[i]) if len(pred_boxes[i]) >= 3*b else len(pred_boxes[i])\n",
    "                if len(gt_boxes[i]) == 0:\n",
    "                    ind = 3\n",
    "                for pbox in pred_boxes[i][:ind]:\n",
    "                    conf = pbox[4]\n",
    "                    if conf > self.kwargs['conf_thresh']:\n",
    "                        col = 'r'\n",
    "                    else:\n",
    "                        col = 'b'\n",
    "                    self.add_bbox(sp, pbox, color=col)\n",
    "         \n",
    "        \n",
    "        box_dir = join(self.kwargs['save_path'], name + \"_boxes\")\n",
    "        makedir_if_not_there(box_dir)\n",
    "        \n",
    "\n",
    "        plt.savefig(\"%s/%s_epoch_%i_boxes_%i.png\"%(box_dir, name, self.epoch, num ))\n",
    "        self.save_out_boxes(box_dir, \"ten_most_conf_boxes_epoch_%i\"%(self.epoch), pred_boxes[0], gt_boxes[0])\n",
    "                \n",
    "        \n",
    "        plt.savefig(\"%s/%s_boxes_%i.png\"%(box_dir, name, num))\n",
    "        self.save_out_boxes(box_dir, \"ten_most_conf_boxes\", pred_boxes[0], gt_boxes[0])\n",
    "        \n",
    "        #plt.show()\n",
    "        plt.clf()\n",
    "\n",
    "    \n",
    "    def save_out_boxes(self,box_dir, name, pred_boxes, gt_boxes):\n",
    "        pred_boxes.sort(lambda a,b: -1 if a[4] > b[4] else 1) \n",
    "        with open(join(box_dir, name), \"a\") as f:\n",
    "            ind = 10 if len(pred_boxes) >= 10 else len(pred_boxes)\n",
    "            for box in pred_boxes[:ind]:\n",
    "                f.write(\"\\n\" + str(box) + \"\\n\")\n",
    "            ind = ind = 10 if len(gt_boxes) >= 10 else len(gt_boxes)\n",
    "            f.write(\"\\n ooooh gt boxes: \\n\")\n",
    "            for box in gt_boxes[:ind]:\n",
    "                f.write(str(box) + \"\\n\")\n",
    "            f.write(\"\\n\\n\\n\")\n",
    "            \n",
    "\n",
    "    def add_bbox(self, subplot, bbox, color):\n",
    "        #box of form center x,y  w,h\n",
    "        x,y,w,h, conf1, cls = bbox\n",
    "        subplot.add_patch(patches.Rectangle(\n",
    "        xy=( y - h / 2., (self.xdim - x) - w / 2.),\n",
    "        width=h,\n",
    "        height=w, lw=2,\n",
    "        fill=False, color=color))\n",
    "        subplot.text(y - h / 2., (self.xdim - x) - w / 2.,self.classes[cls], fontdict={\"color\":color})\n",
    "\n",
    "                \n",
    "#     def plot_reconstructed_ims(self,x,xrec):\n",
    "\n",
    "            \n",
    "    def _plot_reconstructed_ims(self, orig, rec, num=1, name=\"tr\"):\n",
    "        \"\"\"orig and rec is a N,vars,X,y tensor\"\"\"\n",
    "        \"\"\" for 3D they are an N,vars, time_steps,x,y \"\"\"\n",
    "        if self.kwargs[\"3D\"]:\n",
    "            n_ims = orig.shape[2]\n",
    "        else:\n",
    "            n_ims = orig.shape[0]\n",
    "        tmq_ind =6 \n",
    "        psl_ind = 2\n",
    "        channels = [tmq_ind, psl_ind] #orig.shape[1] #ims.shape[1]\n",
    "        plt.figure(4, figsize=(40,40))\n",
    "        plt.clf()\n",
    "        \n",
    "\n",
    "        count=0\n",
    "        for j in channels:\n",
    "            for i in range(n_ims):\n",
    "            \n",
    "                if self.kwargs[\"3D\"]:\n",
    "                    im_slice = slice(0,i,j)\n",
    "                else:\n",
    "                    im_slice = slice(i,j)\n",
    "                count += 1\n",
    "                sp = plt.subplot(n_ims*2,len(channels), count)\n",
    "                if self.kwargs[\"3D\"]:\n",
    "                    sp.imshow(orig[0,j, i][::-1])\n",
    "                else:\n",
    "                    sp.imshow(orig[i,j][::-1])\n",
    "                    \n",
    "                count +=1\n",
    "                sp = plt.subplot(n_ims*2,len(channels), count)\n",
    "                if self.kwargs[\"3D\"]:\n",
    "                    sp.imshow(rec[0,j, i][::-1])\n",
    "                else:\n",
    "                    sp.imshow(rec[i,j][::-1])\n",
    "        \n",
    "        rec_dir = join(self.kwargs['save_path'], name + \"_rec\")\n",
    "        makedir_if_not_there(rec_dir)\n",
    "        \n",
    "        \n",
    "        if self.epoch % 4 == 0:           \n",
    "            plt.savefig(\"%s/%s_epoch_%i_rec_%i.png\"%(rec_dir, name, self.epoch, num))\n",
    "        plt.savefig(\"%s/%s_rec_%i.png\"%(rec_dir, name, num))\n",
    "        #plt.show()\n",
    "        plt.clf()\n",
    "        "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
